#!/bin/bash
# 增强版文件/目录加密工具：支持目录递归、保留结构、跳过已处理文件

usage() {
    echo -e "\033[1;37mUsage: $0 <mode> <source> <output> [password_file]\033[0m"
    echo "  mode:         encrypt 或 decrypt"
    echo "  source:       加密时为目录；解密时为目录（.enc后缀）"
    echo "  output:       加密输出路径；解密输出路径"
    echo "  password_file: 密码文件路径（推荐）"
    echo -e "\n\033[33m示例:"
    echo "  # 加密整个目录（跳过已加密文件）"
    echo "  $0 encrypt ~/documents ~/encrypted_docs ~/.secure_pass"
    echo "  # 解密整个目录（跳过已解密文件）"
    echo "  $0 decrypt ~/encrypted_docs ~/restored_docs"
    echo -e "\033[0m"
    exit 1
}

# 检查参数
if [ $# -lt 4 ]; then usage; fi

MODE=$1
SOURCE=$2
OUTPUT=$3
PASS_FILE=${4:-}

# 核心改进：检查目标文件是否存在
check_target_exists() {
    local src_path="$1"
    local target_path="$2"
    
    if [ -e "$target_path" ]; then
        echo -e "\033[33m⚠ 跳过 [$src_path] -> 目标已存在: $target_path\033[0m"
        return 1
    fi
    return 0
}

# 加密单个文件（增加存在性检查）
encrypt_file() {
    local src_file="$1"
    local out_file="$2"
    local pass_opt="$3"

    # 跳过已存在的加密文件
    if ! check_target_exists "$src_file" "$out_file"; then return; fi
    
    echo -e "\033[32m▶ 加密文件: $src_file -> $out_file\033[0m"
    openssl enc -aes-256-cbc -salt -pbkdf2 $pass_opt -in "$src_file" -out "$out_file"
    [ $? -eq 0 ] && echo -e "\033[1;34m✔ 加密成功\033[0m" || return 1
}

# 解密单个文件（增加存在性检查）
decrypt_file() {
    local src_file="$1"
    local out_file="$2"
    local pass_opt="$3"

    # 跳过已存在的解密文件
    if ! check_target_exists "$src_file" "$out_file"; then return; fi
    
    echo -e "\033[32m▶ 解密文件: $src_file -> $out_file\033[0m"
    openssl enc -d -aes-256-cbc -pbkdf2 $pass_opt -in "$src_file" -out "$out_file"
    if [ $? -eq 0 ]; then
        echo -e "\033[1;34m✔ 解密成功\033[0m"
    else
	echo -e "\033[31m✘ 解密失败\033[0m"
	exit 1
    fi
}

# 递归处理目录（核心逻辑优化）
process_dir() {
    local src_dir="$1"
    local out_dir="$2"
    local pass_opt="$3"
    local mode="$4"
    
    local processed=0 skipped=0 errors=0
    
    while IFS= read -r -d $'\0' file; do
        local rel_path="${file#$src_dir/}"
        local target_path="$out_dir/${rel_path}"
        
        case $mode in
            encrypt)
                target_path="${target_path}.enc"
                if check_target_exists "$file" "$target_path"; then
                    mkdir -p "$(dirname "$target_path")"
		    if encrypt_file "$file" "$target_path" "$pass_opt"; then
                        ((processed++))  # 仅成功时计数
                    else
                        ((errors++))     # 仅失败时计数
                    fi
                else
                    ((skipped++))
                fi
                ;;
            decrypt)
                if [[ "$file" != *.enc ]]; then
                    echo -e "\033[33m⚠ 跳过非加密文件: $file\033[0m"
                    ((skipped++))
                    continue
                fi
                target_path="${target_path%.enc}"
                if check_target_exists "$file" "$target_path"; then
                    mkdir -p "$(dirname "$target_path")"

		    if decrypt_file "$file" "$target_path" "$pass_opt"; then
                        ((processed++))  # 仅成功时计数
                    else
                        ((errors++))     # 仅失败时计数
                    fi
                else
                    ((skipped++))
                fi
                ;;
        esac
    done < <(find "$src_dir" -type f -print0)
    
    # 统计报告
    echo -e "\n\033[1;35m≡ 操作完成: $processed 成功 | $skipped 跳过 | $errors 失败\033[0m"
}

# 密码处理
if [ -n "$PASS_FILE" ]; then
    [ -f "$PASS_FILE" ] || { echo -e "\033[31m✘ 密码文件不存在 [$PASS_FILE]\033[0m"; exit 2; }
    PASS_OPT="-pass file:$PASS_FILE"
fi

# 主流程
case $MODE in
    encrypt)
        if [ -f "$SOURCE" ]; then
	    usage
        elif [ -d "$SOURCE" ]; then
            [ -z "$OUTPUT" ] && { echo -e "\033[31m✘ 加密目录必须指定输出目录\033[0m"; exit 3; }
            process_dir "$SOURCE" "$OUTPUT" "$PASS_OPT" "encrypt"
        else
            echo -e "\033[31m✘ 无效路径: '$SOURCE' (非文件/目录)\033[0m"; exit 4
        fi
        ;;
    decrypt)
        if [ -f "$SOURCE" ]; then
	    usage
        elif [ -d "$SOURCE" ]; then
            [ -z "$OUTPUT" ] && { echo -e "\033[31m✘ 解密目录必须指定输出目录\033[0m"; exit 3; }
            process_dir "$SOURCE" "$OUTPUT" "$PASS_OPT" "decrypt"
        else
            echo -e "\033[31m✘ 无效路径: '$SOURCE' (非文件/目录)\033[0m"; exit 4
        fi
        ;;
    *) echo -e "\033[31m✘ 无效模式: [$MODE]\033[0m"; usage ;;
esac
